
import cntk as C

import onnxruntime
import numpy as np

def convet():
    model_path = 'checkpoint/danbooru-resnet_custom_v1-p3/model.cntk'
    ctnk_model = C.load_model(model_path)
    ctnk_model.save('checkpoint/model.onnx', format=C.ModelFormat.ONNX)
    return ctnk_model



def test_cntk(ctnk_model):
    ort_session = onnxruntime.InferenceSession("checkpoint/model.onnx")

    x = np.random.rand(1, 3, 299, 299).astype(np.float32)
    # compute ONNX Runtime output prediction
    ort_inputs = {ort_session.get_inputs()[0].name: x}
    ort_outs = ort_session.run(None, ort_inputs)

    # compute the ctnk output
    cntk_out = ctnk_model.eval(x)
    np.testing.assert_allclose(np.array(cntk_out), ort_outs[0], rtol=1e-03, atol=1e-05)


if __name__ == '__main__':
    ctnk_model=convet()
    test_cntk(ctnk_model)